import request from './request'
import { API_CONFIG } from './config'

class WebSocketClient {
  constructor() {
    this.ws = null
    this.isConnected = false
    this.messageHandlers = new Set()
    this.errorHandlers = new Set()
    this.metadataHandlers = new Set()
    this.completionHandlers = new Set()
  }

  connect() {
    return new Promise((resolve, reject) => {
      this.ws = new WebSocket(API_CONFIG.PREDICTION.WS_URL)

      this.ws.onopen = () => {
        this.isConnected = true
        console.log('WebSocket connected')
        resolve()
      }

      this.ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data)

          // 根据消息类型分发
          if (data.type === 'metadata') {
            this.metadataHandlers.forEach(handler => handler(data))
          } else if (data.status === 'completed') {
            this.completionHandlers.forEach(handler => handler(data))
          } else if (data.error) {
            this.errorHandlers.forEach(handler => handler(data.error))
          } else {
            // 预测数据
            this.messageHandlers.forEach(handler => handler(data))
          }
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error)
        }
      }

      this.ws.onerror = (error) => {
        console.error('WebSocket error:', error)
        this.errorHandlers.forEach(handler => handler(error))
        reject(error)
      }

      this.ws.onclose = () => {
        this.isConnected = false
        console.log('WebSocket disconnected')
      }
    })
  }

  disconnect() {
    if (this.ws && this.isConnected) {
      this.ws.close()
    }
  }

  onMessage(handler) {
    this.messageHandlers.add(handler)
    return () => this.messageHandlers.delete(handler)
  }

  onMetadata(handler) {
    this.metadataHandlers.add(handler)
    return () => this.metadataHandlers.delete(handler)
  }

  onCompletion(handler) {
    this.completionHandlers.add(handler)
    return () => this.completionHandlers.delete(handler)
  }

  onError(handler) {
    this.errorHandlers.add(handler)
    return () => this.errorHandlers.delete(handler)
  }

  sendMessage(data) {
    if (this.ws && this.isConnected) {
      this.ws.send(typeof data === 'string' ? data : JSON.stringify(data))
    } else {
      throw new Error('WebSocket is not connected')
    }
  }

  stopPrediction() {
    if (this.ws && this.isConnected) {
      this.ws.send('stop')
    }
  }
}

export const predictionApi = {
  /**
   * 创建WebSocket连接
   * @returns {WebSocketClient} WebSocket客户端实例
   */
  createWebSocket() {
    return new WebSocketClient()
  },

  /**
   * 开始预测
   * @param {WebSocketClient} wsClient - WebSocket客户端实例
   * @param {Object} config - 预测配置
   * @param {string} config.data_file - 数据文件名
   * @param {number} config.seq_len - 序列长度
   * @param {number} config.pred_len - 预测长度
   * @param {boolean} config.use_gpu - 是否使用GPU
   * @param {string} config.export_format - 导出格式
   */
  startPrediction(wsClient, config) {
    wsClient.sendMessage({
      data_file: config.data_file,
      seq_len: config.seq_len || 126,
      pred_len: config.pred_len || 4,
      use_gpu: config.use_gpu !== undefined ? config.use_gpu : true,
      export_format: config.export_format || API_CONFIG.PREDICTION.EXPORT_FORMAT
    })
  },

  /**
   * 停止预测
   * @param {WebSocketClient} wsClient - WebSocket客户端实例
   */
  stopPrediction(wsClient) {
    wsClient.stopPrediction()
  },

  /**
   * 导出预测结果为CSV
   * @returns {Promise} 返回文件下载响应
   */
  exportPredictionCSV() {
    return request({
      url: API_CONFIG.PREDICTION.EXPORT_CSV,
      method: 'GET',
      responseType: 'blob'
    })
  }
} 